import socket
import struct
import time
from collections import defaultdict

class EncoderDataParser:
    def __init__(self):
        # 存储每个轴的最新状态
        self.axis_states = defaultdict(lambda: {
            'timestamp': 0,
            'axis_count': 0,
            'stepper_count': 0
        })
        
        # 数据包格式定义
        self.PACKET_FORMATS = {
            'incremental': struct.Struct('<B b b b'),  # 4字节包: [类型, dt, axis_delta, stepper_delta]
            'full': struct.Struct('<B q i i')          # 17字节包: [类型, timestamp, axis_count, stepper_count]
        }
        
        # 缓冲区用于存储不完整的数据
        self.buffer = bytearray()
        
    def parse_packet(self, data):
        """解析接收到的数据包"""
        results = []
        self.buffer.extend(data)
        
        while len(self.buffer) >= 4:  # 最小包是4字节
            pkt_type = self.buffer[0]
            axis_id = pkt_type & 0x7F  # 提取轴ID (低7位)
            is_full_packet = (pkt_type & 0x80) != 0  # 检查最高位
            
            if is_full_packet:
                # 完整包需要17字节
                if len(self.buffer) < 17:
                    break  # 等待更多数据
                
                # 解析完整包
                packet = self.PACKET_FORMATS['full'].unpack_from(self.buffer[:17])
                _, timestamp, axis_count, stepper_count = packet
                
                # 更新状态
                self.axis_states[axis_id]['timestamp'] = timestamp
                self.axis_states[axis_id]['axis_count'] = axis_count
                self.axis_states[axis_id]['stepper_count'] = stepper_count
                
                # 添加到结果
                results.append({
                    'axis': axis_id,
                    'type': 'full',
                    'timestamp': timestamp,
                    'axis_count': axis_count,
                    'stepper_count': stepper_count
                })
                
                # 从缓冲区移除已处理的数据
                self.buffer = self.buffer[17:]
                
            else:
                # 增量包只需4字节
                packet = self.PACKET_FORMATS['incremental'].unpack_from(self.buffer[:4])
                _, dt, axis_delta, stepper_delta = packet
                
                # 获取当前状态
                state = self.axis_states[axis_id]
                
                # 更新状态 (转换为有符号数)
                axis_delta = self._to_signed(axis_delta)
                stepper_delta = self._to_signed(stepper_delta)
                
                state['timestamp'] += dt
                state['axis_count'] += axis_delta
                state['stepper_count'] += stepper_delta
                
                # 添加到结果
                results.append({
                    'axis': axis_id,
                    'type': 'incremental',
                    'dt': dt,
                    'axis_delta': axis_delta,
                    'stepper_delta': stepper_delta,
                    'axis_count': state['axis_count'],
                    'stepper_count': state['stepper_count']
                })
                
                # 从缓冲区移除已处理的数据
                self.buffer = self.buffer[4:]
        
        return results
    
    def _to_signed(self, value):
        """将8位有符号数转换为正确的Python有符号整数"""
        return value - 256 if value > 127 else value

def tcp_client(host: str, port: int, timeout=5.0):
    """
    TCP 客户端，支持解析编码器数据包
    :param host: 远程IP地址
    :param port: 远程端口
    :param timeout: 超时时间(秒)
    """
    parser = EncoderDataParser()
    
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        f = None
        try:
            s.settimeout(timeout)
            s.connect((host, port))
            print(f"Connected to {host}:{port}")
            f = open(f"enc_{time.time()}.enc","wb+")
            s.send(b"Hello!")
            s.setblocking(True)
            # 持续接收数据
            while True:
                data = s.recv(1024)  # 每次最多接收1024字节
                if not data:
                    print("Connection closed by server")
                    break
                f.write(data)
        except socket.timeout:
            print("Connection timeout")
        except ConnectionRefusedError:
            print("Connection refused (server not running)")
        except KeyboardInterrupt:
            print("Disconnected by user")
        except Exception as e:
            print(f"Error: {str(e)}")
        finally:
            print("Closing connection")
        if f:
            f.close()

# 使用示例
if __name__ == "__main__":
    # 替换为你的ESP32 IP地址
    tcp_client("192.168.0.118", 12345)